import _thread
import time

"""
线程间共享数据，通过 lock 保证数据安全

用法：
1. 创建子类，继承 ThreadSharedData，子类中定义自己的数据
2. 访问数据时，使用 with 语句
"""
class ThreadSharedData:
    def __init__(self):
        self.lock = _thread.allocate_lock()
        
    def __enter__(self):
        self.lock.acquire()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.lock.release()

"""
该类继承于 ThreadSharedData，主要用于一主一子线程，主线程可请求子线程退出。
    1. 主线程通过 request_and_wait_thread_exit() 请求子线程退出；
    2. 子线程通过不停检查 check_thread_exit_request() 检查是否收到主线程发来的退出请求。

用户也可继续继承本类，增加更多自定义的线程间共享数据。
"""
class ThreadSharedData_StopOnly(ThreadSharedData):
    def __init__(self):
        super().__init__()
        self._request_exit_flag = False
        self._thread_running = False

    """主线程中调用，当创建子线程成功时调用该函数"""
    def on_thread_start(self):
        with self:
            self._request_exit_flag = False
            self._thread_running = True

    """主线程中调用，请求子线程退出，并等待退出。注意，这是一个阻塞调用"""
    def request_and_wait_thread_exit(self):
        with self:
            self._request_exit_flag = True

        while True:
            with self:
                if not self._thread_running:
                    return
                
            time.sleep(0.1)

    """子线程中调用，检查是否收到主线程发来的退出请求。如果检查到要退出，还可以执行 on_exit_callback 回调函数"""
    def check_thread_exit_request(self, on_exit_callback=None):
        should_exit = False

        with self:
            if self._request_exit_flag:
                should_exit = True

        if should_exit:
            if on_exit_callback is not None:
                on_exit_callback()

            with self:
                self._thread_running = False

            _thread.exit()

    """子线程也可主动退出，调用该函数"""
    def quit_thread(self):
        with self:
            self._thread_running = False

        _thread.exit()



"""
线程间通讯压力测试代码
"""
def pressure_test():
    class Test:
        NUM_DATA = 100

        def __init__(self):
            class MyThreadSharedData(ThreadSharedData_StopOnly):
                def __init__(self):
                    super().__init__()

                    from collections import deque
                    self.my_queue: deque = deque((), Test.NUM_DATA)
            
            """创建线程间共享数据"""
            self.shared_data = MyThreadSharedData()

        """这是生产者线程"""
        def producer(self):
            for i in range(Test.NUM_DATA):
                with self.shared_data: # with 语句保证数据安全
                    self.shared_data.my_queue.append(i)

                time.sleep(0.001)

            # 生产完成后，等待主线程发来退出请求
            while True:
                self.shared_data.check_thread_exit_request()
                time.sleep(1)

        """这是运行于主线程上的消费者"""
        def consumer(self):
            counter = 0
            while True:
                with self.shared_data: # with 语句保证数据安全
                    if len(self.shared_data.my_queue) > 0:
                        data = self.shared_data.my_queue.pop()
                        counter += 1
                        # print(data)

                if counter == Test.NUM_DATA:
                    break

                time.sleep(0.002)

            print('All data received, requesting thread exit...')
            self.shared_data.request_and_wait_thread_exit()
            print('Thread exited.')

        def do_test(self):
            _thread.start_new_thread(self.producer, ())
            self.consumer()

    import gc
    for i in range(1000):
        print(f'Running the {i}th test...')
        test = Test()

        try:
            test.do_test()
            print('Test passed.')
            print(f'-----------Memory free: {gc.mem_free()}----------------')
        except OSError as e:
            print(e)
            print('Test Failed.')
            
            gc.collect()
            for _ in range(15):
                print(f'-----------Memory free: {gc.mem_free()}----------------')

                time.sleep(1)

if __name__ == '__main__':
    pressure_test()